#!/bin/bash

# 多模态MoE分类器训练脚本
# 使用方法: ./train.sh [GPU_ID] [NUM_CLASSES] [BATCH_SIZE] [EPOCHS]

# 设置默认参数
GPU_ID=${1:-2}           # 默认使用GPU 2
NUM_CLASSES=${2:-89}    # 默认89个类别
BATCH_SIZE=${3:-128}      # 默认批次大小128
EPOCHS=${4:-30}          # 默认训练30轮

# 数据路径配置
# TRAIN_CSV="/home74/liguangzhen/Paper/Disease/MoE/MoE_classifier/data_bp/train_fmt.csv"
# VAL_CSV="/home74/liguangzhen/Paper/Disease/MoE/MoE_classifier/data_bp/val_fmt.csv"
# IMAGE_ROOT="/repository7403/liguangzhen/Disease/bp_18000/image"
TRAIN_CSV="/home74/liguangzhen/Paper/Disease/MoE/MoE_classifier/data_sclod/data/train_fmt.csv"
VAL_CSV="/home74/liguangzhen/Paper/Disease/MoE/MoE_classifier/data_sclod/data/val_fmt.csv"
IMAGE_ROOT="/repository7403/liguangzhen/Disease/scold_12000/images"
OUTPUT_DIR="runs_cls"

# 模型配置
IMG_BACKBONE="resnet50"
TEXT_MODEL="distilbert-base-uncased"
FEATURE_DIM=512
NUM_EXPERTS=3
TOP_K=2
# 新策略默认值（可通过环境或参数覆盖）
VISUAL_PRETRAINED=1
USE_VISUAL_AUX=1
VISUAL_AUX_WEIGHT=0.4
GATING_STRATEGY="prob_norm"  # 或 softmax_topk/softmax_temp/power
GATING_TEMPERATURE=1.0
GATING_POWER_ALPHA=1.0
SECOND_PROB_THRESHOLD=0.05
AUGMENT_TRAIN="strong"
RAND_MAGNITUDE=7
JITTER=0.3
GATE_MONITOR=1
VISUAL_MANDATORY_EPOCHS=5
VISUAL_MANDATORY_INIT_STRENGTH=0.5

# 训练配置
LR=3e-4
WEIGHT_DECAY=1e-4
IMAGE_SIZE=224
MAX_LEN=64
LB_COEFF=0.01
NUM_WORKERS=0  # 设置为0避免多进程问题

echo "=========================================="
echo "🚀 多模态MoE分类器训练"
echo "=========================================="
echo "📊 GPU ID: $GPU_ID"
echo "🏷️  类别数: $NUM_CLASSES"
echo "📦 批次大小: $BATCH_SIZE"
echo "🔄 训练轮数: $EPOCHS"
echo "📁 训练数据: $TRAIN_CSV"
echo "📁 验证数据: $VAL_CSV"
echo "🖼️  图像路径: $IMAGE_ROOT"
echo "💾 输出目录: $OUTPUT_DIR"
echo "=========================================="

# 检查数据文件是否存在
if [ ! -f "$TRAIN_CSV" ]; then
    echo "❌ 错误: 训练数据文件不存在: $TRAIN_CSV"
    exit 1
fi

if [ ! -f "$VAL_CSV" ]; then
    echo "❌ 错误: 验证数据文件不存在: $VAL_CSV"
    exit 1
fi

if [ ! -d "$IMAGE_ROOT" ]; then
    echo "❌ 错误: 图像目录不存在: $IMAGE_ROOT"
    exit 1
fi

# 创建输出目录
mkdir -p $OUTPUT_DIR

# 运行训练
python train_classifier.py \
    --train_csv $TRAIN_CSV \
    --val_csv $VAL_CSV \
    --image_root $IMAGE_ROOT \
    --num_classes $NUM_CLASSES \
    --img_backbone $IMG_BACKBONE \
    --text_model $TEXT_MODEL \
    --feature_dim $FEATURE_DIM \
    --num_experts $NUM_EXPERTS \
    --top_k $TOP_K \
    --batch_size $BATCH_SIZE \
    --epochs $EPOCHS \
    --lr $LR \
    --weight_decay $WEIGHT_DECAY \
    --image_size $IMAGE_SIZE \
    --max_len $MAX_LEN \
    --lb_coeff $LB_COEFF \
    --output_dir $OUTPUT_DIR \
    --gpu_id $GPU_ID \
    --num_workers $NUM_WORKERS \
    --augment_train $AUGMENT_TRAIN \
    --rand_magnitude $RAND_MAGNITUDE \
    --jitter $JITTER \
    $( [ "$VISUAL_PRETRAINED" = "1" ] && echo "--visual_pretrained" ) \
    $( [ "$USE_VISUAL_AUX" = "1" ] && echo "--use_visual_aux" ) \
    --visual_aux_weight $VISUAL_AUX_WEIGHT \
    --gating_strategy $GATING_STRATEGY \
    --gating_temperature $GATING_TEMPERATURE \
    --gating_power_alpha $GATING_POWER_ALPHA \
    --second_prob_threshold $SECOND_PROB_THRESHOLD \
    $( [ "$GATE_MONITOR" = "1" ] && echo "--gate_monitor" ) \
    --visual_mandatory_epochs $VISUAL_MANDATORY_EPOCHS \
    --visual_mandatory_init_strength $VISUAL_MANDATORY_INIT_STRENGTH

echo "=========================================="
echo "🎉 训练脚本执行完成!"
echo "=========================================="